'''
Inference for Composition-1k Dataset.

Run:
python inference.py \
    --config-dir path/to/config
    --checkpoint-dir path/to/checkpoint
    --inference-dir path/to/inference
    --data-dir path/to/data
'''
import os
import torch
import torch as th
from PIL import Image
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import functional as F
from os.path import join as opj
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.config import LazyConfig, instantiate
from detectron2.engine import default_argument_parser

import warnings
import time

warnings.filterwarnings('ignore')


def timeit_fps(repeat_times=5, bs=1):
    def timeit_decorator(func):
        def timeit_wrapper(*args, **kwargs):
            # warm up
            print(f'Warm up {func.__name__} started...')
            result = func(*args, **kwargs)
            print(f'Warm up {func.__name__} done')

            start_time = time.perf_counter()
            for _ in range(repeat_times):
                result = func(*args, **kwargs)
                assert result is not None
            end_time = time.perf_counter()
            total_time = end_time - start_time

            total_time /= repeat_times
            total_time /= bs

            print(f'Function {func.__name__}{args} bs={bs}, {kwargs}:\n'
                  f'\tTook {total_time * 1000:.4f} ms\n'
                  f'\tFPS: {1 / total_time:.2f} /s\n')
            return result

        return timeit_wrapper

    return timeit_decorator


def measure_fps(config_dir='configs/ViTMatte_B_100ep.py',
                img_size: int = 1920, bs: int = 1):
    # initializing model
    cfg = LazyConfig.load(config_dir)
    model = instantiate(cfg.model)
    model.cuda()
    model.eval()

    image = th.randn(bs, 3, img_size, img_size, device='cuda')
    trimap = th.randn(bs, 1, img_size, img_size, device='cuda')

    batch = {'image': image, 'trimap': trimap, 'image_name': 'random_test_image'}

    @timeit_fps(bs=bs)
    def f(model=model, batch=batch):
        with torch.inference_mode():
            out = model(batch)['phas'].flatten(0, 2)
        # print(len(out))
        return out

    f()


if __name__ == '__main__':
    # compare_generator()
    config_dirs = [
        # 'configs/ViTMatte_S_100ep.py',
        'configs/ViTMatte_B_100ep.py'
    ]
    img_sizes = [736, 1024, 1216, 1312, 1408, 1504, 1600, 1728, 1824, 2016, 2208, 2400]


    bss = [1] * len(img_sizes)

    for img_size, bs in zip(img_sizes, bss):
        for config_dir in config_dirs:
            print(f'cfg={config_dir}, img_size={img_size}, bs={bs}')
            measure_fps(config_dir=config_dir, img_size=img_size, bs=bs)
